package main

func subarraySum2(nums []int, k int) int {
	preSum := make([]int, len(nums)+1)
	for i := 1; i <= len(nums); i++ {
		preSum[i] = preSum[i-1] + nums[i-1]
	}
	count := 0
	for i := 0; i < len(preSum); i++ {
		for j := i + 1; j < len(preSum); j++ {
			sum := preSum[j] - preSum[i]
			if sum == k {
				count++
			}
		}
	}
	return count
}

func subarraySum(nums []int, k int) int {
	preSum := make([]int, len(nums)+1)
	preSumCount := make(map[int]int) // presum -> count
	preSumCount[0] = 1
	count := 0
	for i := 1; i <= len(nums); i++ {
		preSum[i] = preSum[i-1] + nums[i-1]
		n := preSumCount[preSum[i] - k]
		count += n
		preSumCount[preSum[i]] = preSumCount[preSum[i]]  + 1
	}
	return count
}

func main() {
	ans := subarraySum2([]int{1, 1, 1}, 2)
	println(ans)
}
